from Transformer_config import *  # 导入配置文件中的模型参数和全局变量

# Embeddings类 实现思路分析
# 1 init函数 (self, d_model, vocab)
    # 设置类属性 定义词嵌入层 self.lut层
# 2 forward(x)函数
    # self.lut(x) * math.sqrt(self.d_model)
class Embedding(nn.Module):
    """
    词嵌入层，将离散的词汇索引转换为连续的向量表示

    参数:
        vocabulary_size (int): 词汇表大小，表示可处理的不同词汇数量
        embedding_dim (int): 词嵌入维度，每个词汇对应的特征向量长度
    """
    def __init__(self,vocabulary_size,embedding_dim):
        # 参数vocab   词汇表大小
        # 参数d_model 每个词汇的特征尺寸 词嵌入维度
        super().__init__()  # 调用nn.Module基类的初始化方法
        self.vocabulary_size = vocabulary_size  # 存储词汇表大小
        self.embedding_dim = embedding_dim  # 存储词嵌入维度

        # 定义吃嵌入层
        # 创建PyTorch嵌入层：vocabulary_size表示词汇表大小，embedding_dim表示向量维度
        self.embedding_layer = nn.Embedding(vocabulary_size,embedding_dim,device=device)

    def forward(self,x):
        """
        前向传播过程：将输入索引转换为嵌入向量并缩放

        参数:
            x (Tensor): 输入词汇索引张量，形状为[batch_size, seq_len]

        返回:
            Tensor: 缩放后的词嵌入向量，形状为[batch_size, seq_len, embedding_dim]
        """
        # 将x传给self.lut并与根号下self.d_model相乘作为结果返回
        # x经过词嵌入后 增大x的值, 词嵌入后的embedding_vector+位置编码信息,值量纲差差不多
        # 通过嵌入层将输入索引转换为稠密向量
        x = self.embedding_layer(x)
        # 缩放嵌入向量使其与位置编码的量级匹配
        x = x * math.sqrt(self.embedding_dim)
        return x  # 返回缩放后的词嵌入结果

# 位置编码器类PositionalEncoding 实现思路分析
# 1 init函数  (self, d_model, dropout, max_len=5000)
#   super()函数 定义层self.dropout
#   定义位置编码矩阵pe  定义位置列-矩阵position 定义变化矩阵div_term
#   套公式div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0)/d_model))
#   位置列-矩阵 * 变化矩阵 阿达码积my_matmulres
#   给pe矩阵偶数列奇数列赋值 pe[:, 0::2] pe[:, 1::2]
#   pe矩阵注册到模型缓冲区 pe.unsqueeze(0)三维 self.register_buffer('pe', pe)
# 2 forward(self, x) 返回self.dropout(x)
#   给x数据添加位置特征信息 x = x + Variable( self.pe[:,:x.size()[1]], requires_grad=False)
class PositionalEncoding(nn.Module):
    """
    位置编码层，为序列添加位置信息以保留顺序关系

    参数:
        embedding_dim (int): 词嵌入维度，需与词嵌入层输出维度一致
        max_len (int): 支持的最大序列长度
        dropout (float): Dropout层的丢弃概率
    """
    def __init__(self,embedding_dim,max_len,dropout):
        # 参数d_model 词嵌入维度 eg: 512个特征
        # 参数max_len 单词token个数 eg: 64个单词
        super().__init__()  # 调用nn.Module基类的初始化方法

        # 定义dropout层
        # 创建Dropout层，p=dropout指定丢弃概率
        self.dropout = nn.Dropout(p=dropout)

        # 思路：位置编码矩阵 + 特征矩阵 相当于给特征增加了位置信息
        # 定义位置编码矩阵PE eg pe[60, 512], 位置编码矩阵和特征矩阵形状是一样的
        # 初始化位置编码矩阵：max_len×embedding_dim的全零张量
        pe = torch.zeros(max_len,embedding_dim,device=device)

        # 定义位置列-矩阵position  数据形状[max_len,1] eg: [0,1,2,3,4...64]^T
        # 创建位置索引序列并增加一个维度变为列向量
        position = torch.arange(0,max_len).unsqueeze(1)

        # 定义变化矩阵div_term [1,256]
        # torch.arange(start=1, end=512, 2)结果并不包含end。在start和end之间做一个等差数组 [0, 2, 4, 6 ... 510]
        # 计算频率调节因子：用于生成不同频率的正弦/余弦波
        div_term = torch.exp(torch.arange(0,embedding_dim,2) * -(math.log(10000.0)/embedding_dim))

        # 位置列-矩阵 @ 变化矩阵 做矩阵运算 [64*1]@ [1*256] ==> 64 *256
        # 矩阵相乘也就是行列对应位置相乘再相加，其含义，给每一个列属性（列特征）增加位置编码信息
        # 计算位置和频率的乘积矩阵
        my_matmulres = position * div_term

        # 给位置编码矩阵奇数列，赋值sin曲线特征
        # 对偶数索引位置应用正弦函数
        pe[:,0::2] = torch.sin(my_matmulres)
        # 给位置编码矩阵偶数列，赋值cos曲线特征
        # 对奇数索引位置应用余弦函数
        pe[:,1::2] = torch.cos(my_matmulres)

        # 形状变化 [60,512]-->[1,64,512]
        # 增加批次维度：从[seq_len, d_model]变为[1, seq_len, d_model]
        pe = pe.unsqueeze(0)

        # 把pe位置编码矩阵 注册成模型的持久缓冲区buffer; 模型保存再加载时，可以根模型参数一样，一同被加载
        # 什么是buffer: 对模型效果有帮助的，但是却不是模型结构中超参数或者参数，不参与模型训练
        # 将位置编码注册为模型缓冲区（不参与训练但会保存）
        self.register_buffer('pe',pe)

    def forward(self,x):
        """
        前向传播过程：为输入序列添加位置编码信息

        参数:
            x (Tensor): 输入张量，形状为[batch_size, seq_len, embedding_dim]

        返回:
            Tensor: 添加位置编码后的张量，形状不变
        """
        # x--》来自于embedding之后的结果--》[batch_size, seq_len, embed_dim]-->[2, 4, 512]
        # 将x和位置编码的信息进行融合
        # 添加位置编码：截取与输入序列长度匹配的部分
        x = x + self.pe[:,:x.shape[1]]
        # 应用Dropout防止过拟合
        x = self.dropout(x)
        return x  # 返回添加位置编码后的结果


def untest_input(x):
    """
    测试函数：验证词嵌入和位置编码模块的功能

    参数:
        x (Tensor): 输入词汇索引张量，形状为[batch_size, seq_len]

    返回:
        Tensor: 经过词嵌入和位置编码处理后的张量
    """
    # 创建词嵌入层实例
    my_embedding = Embedding(vocabulary_size,embedding_dim)

    # 应用词嵌入层处理输入
    embed_x = my_embedding(x)

    # 创建位置编码层实例
    my_positiona = PositionalEncoding(embedding_dim,max_len,dropout)
    # 应用位置编码层处理嵌入结果
    pe_result = my_positiona(embed_x)

    # 打印原始输入信息
    print(f'原始输入为：{x}')
    # 打印词嵌入处理结果
    print(f'{"*"*30}词嵌入{"*"*30}')
    print(f'词嵌入后的形状：{embed_x.shape}')
    # 打印位置编码处理结果
    print(f'{"*"*30}位置编码{"*"*30}')
    print(f'位置编码后的形状：{pe_result.shape}')

    # 返回处理后的结果
    return pe_result

if __name__ == '__main__':
    """
    主程序入口：执行测试函数
    """
    # 创建测试输入张量：2个样本，每个样本4个token
    test_input = torch.tensor([[1,2,3,4],[5,6,7,8]])
    # 执行测试函数
    untest_input(test_input)
